import baostock as bs
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import sys
import matplotlib.gridspec as gridspec
import talib
class DefTypesPool():
    def __init__(self):
        self.routes = {}
    # 注册绘图函数
    def route_types(self, types_str):
        def decorator(f):
            self.routes[types_str] = f
            return f
        return decorator
    # 发现绘图函数
    def route_output(self, path):
        function_val = self.routes.get(path)
        if function_val:
            return function_val
        else:
            raise ValueError('Route "{}"" has not been registered'.format(path))
        
class MplTypesDraw():
    mpl = DefTypesPool() 

    @mpl.route_types(u"line")
    def line_plot(df_index, df_dat, graph):
        # 绘制Line图
        for key, val in df_dat.items():
            graph.plot(np.arange(0,len(val)), val, label=key, lw=1.0)

    
    @mpl.route_types(u"bar")
    def bar_plot(df_index, df_dat, graph):
        # 绘制bar图-Volume
        graph.bar(np.arange(0, len(df_index)), df_dat['bar_red'], facecolor='red')
        graph.bar(np.arange(0, len(df_index)), df_dat['bar_green'], facecolor='green')

    @mpl.route_types(u"annotate")
    def line_plot(df_index, df_dat, graph):
        #绘制annotate图
        for key, val in df_dat.items():
            for kl_index, today in val['andata'].iterrows():
                x_posit = df_index.get_loc(kl_index)
                graph.annotate(u"{}\n{}".format(key, today.name.strftime("%m.%d")),
                               xy=(x_posit, today[val['xy_y']]),
                               xycoords='data',
                               xytext=(val['xytext'][0], val['xytext'][1]),
                               va = val['va'], #点在标注下方
                               textcoords='offset points',
                               fontsize=val['fontsize'],
                               arrowprops=val['arrow'])
                
    @mpl.route_types(u"hline")
    def hline_plot(df_index, df_dat, graph):
        #绘制hline图
        for key, val in df_dat.items():
            graph.axhline(val['pos'], c=val['c'],label=key)

    @mpl.route_types(u"filltrade")
    def filltrade_plot(df_index, df_dat, graph):
        # 绘制filltrade 图
        signal_shift = df_dat['signal'].shift(1)
        # 最前面的NAN用-1填充
        signal_shift.fillna(value=-1, inplace=True)
        list_signal = np.sign(df_dat['signal'] - signal_shift)
        bs_signal = list_signal[list_signal != 0]

        skip_days = False
        for k1_index, value in bs_signal.items():
            if (value == 1) and (skip_days == False):
                start = df_index.get_loc(k1_index)
                skip_days = True
            elif (value == -1) and (skip_days == True):
                end = df_index.get_loc(k1_index)
                skip_days = False

                if df_dat['jdval'].iloc[end-1] < df_dat['jdval'].iloc[start]: #赔钱显示绿色
                    graph.fill_between(np.arange(start, end), 0, df_dat['jdval'][start:end],color='green', alpha=0.38)
                    is_win = False
                else: #赚钱显示红色
                    graph.fill_between(np.arange(start, end), 0, df_dat['jdval'][start:end],color='red',alpha=0.38)
                    is_win = True
                graph.annotate('获利\n' if is_win else '亏损\n',
                               xy = (end, df_dat['jdval'].asof(k1_index)),
                               xytext = (df_dat['xytext'][0], df_dat['xytext'][1]),
                               xycoords='data',
                               va = df_dat['va'],#点在标注下方
                               textcoords = 'offset points',
                               fontsize=df_dat['fontsize'],
                               arrowprops=df_dat['arrow']
                               )
                # 整个时间序列填充为底色blue透明度alpha小于后标注区间颜色
                graph.fill_between(np.arange(0, len(df_index)), 0, df_dat['jdval'], color='blue',alpha=.08)

class MultiTraceIf(MplTypesDraw):
    app = DefTypesPool()
    ###########################
    @app.route_types(u"cash_profit") #cash profit and retracement
    def cash_profit_graph(stock_dat, sub_graph, cash_hold = 100000, slippage = 0.01,c_rate = 5.0 / 10000, t_rate = 1.0 / 1000):
        posit_num = 0 #持股数目
        skip_days = False # 持股/持币状态

        # 最大风险回撤-资金最大回撤
        # 绝对收益-资金的度量
        for k1_index, today in stock_dat.iterrows():
            # 买入/卖出执行代码
            if today.Signal == 1 and skip_days == False: # 买入
                skip_days = True
                posit_num = int(cash_hold / (today.Close + slippage)) #将资金转化为股票
                posit_num = int(posit_num / 100) * 100 # 买入股票至少100股，对posit_num向下取整百

                buy_cash = posit_num * (today.Close + slippage) # 计算买入股票所需现金
                # 计算手续费，不足5元按5元收，并保留2位小数
                commission = round(max(buy_cash * c_rate,5),2)
                cash_hold = cash_hold - buy_cash - commission
            elif today.Signal == -1 and skip_days == True: #卖出 避免未买先卖
                skip_days = False
                sell_cash = posit_num * (today.Close - slippage) # 计算卖出股票所得现金
                # 计算手续费,不足5元按5元收，并保留2位小数
                commission = round(max(sell_cash * c_rate,5),2)
                # 计算印花税，保留2位小数
                tax = round(sell_cash * t_rate, 2)
                cash_hold = cash_hold + sell_cash - commission - tax #剩余现金
            if skip_days == True: #持股
                stock_dat.loc[k1_index, 'total'] = posit_num * today.Close + cash_hold
            else: #空仓
                stock_dat.loc[k1_index, 'total'] = cash_hold
            
            # expanding()计算资金曲线当前的滚动最高值
            stock_dat['max_total'] = stock_dat['total'].expanding().max()
            # 计算资金曲线在滚动最高值之后所回撤的百分比
            stock_dat['per_total'] = stock_dat['total'] / stock_dat['max_total']

        min_point_df = stock_dat.sort_values(by = ['per_total'])[0:1]
        min_point_total = min_point_df.per_total
        max_point_df = stock_dat[stock_dat.index <= min_point_total.index[0]].sort_values(by = ['total'], ascending= False)[0:1]
        max_point_total = max_point_df.total

        print("资金最大回撤%5.2f%% 从%s开始至%s结束" % ((1 - min_point_total.values[0]) * 100, max_point_total.index[0], min_point_total.index[0]))
        
        line_total = "资金总体收益 %d; 上涨幅度 %2.f%%" % (stock_dat['total'].iloc[-1], (stock_dat['total'].iloc[-1] - 100000) / 100000 * 100)
        print(line_total)
        max_total = "资金滚动最高值"

        type_dict = {line_total : stock_dat.total, max_total: stock_dat.max_total,}
        view_function = MplTypesDraw.mpl.route_output(u"line")
        view_function(stock_dat.index, type_dict, sub_graph)
        type_dict = {u'资金最大回撤\n{}'.format(1 - min_point_total.values[0]):
                     {
                         'andata':min_point_df,
                         'va':'top',
                         'xy_y': 'total',
                         'xytext': (0, stock_dat['High'].mean()),
                         'fontsize': 8,
                         'arrow': dict(facecolor='green', shrink=0.1)
                     },}
        
        view_function = MplTypesDraw.mpl.route_output(u"annotate")
        view_function(stock_dat.index, type_dict, sub_graph)
    
    @app.route_types(u"cmp_profit")
    def cmp_profit_graph(stock_dat, sub_graph, para_dat):
        #相对收益-策略 vs 基准
        stock_dat['benchmark_profit_log'] = np.log(stock_dat.Close / stock_dat.Close.shift(1))
        stock_dat.loc[stock_dat.Signal == -1,'Signal'] = 0
        stock_dat['trend_profit_log'] = stock_dat['Signal'] * stock_dat.benchmark_profit_log
        line_trend_key = "策略收益%.2f" % stock_dat['trend_profit_log'].cumsum().iloc[-1]
        line_bench_key = "基准收益%.2f" % stock_dat['benchmark_profit_log'].cumsum().iloc[-1]
        print("资金相对收益: %s vs %s" % (line_trend_key, line_bench_key)) 
        type_dict = {line_bench_key: stock_dat['benchmark_profit_log'].cumsum(), line_trend_key: stock_dat['trend_profit_log'].cumsum()}
        view_function = MplTypesDraw.mpl.route_output(u"line")
        view_function(stock_dat.index, type_dict, sub_graph)
    
    @app.route_types(u"close_retrace")
    def close_retrace_graph(stock_dat, sub_graph, para_dat):
        #度量策略最大风险回撤--收盘价最大回撤
        #计算收盘价曲线当前的滚动最高值
        stock_dat['max_close'] = stock_dat['Close'].expanding().max()
        #计算收盘价曲线再滚动最高值之后所回撤的百分比
        stock_dat['per_close'] = stock_dat['Close'] / stock_dat['max_close']

        #找出收盘价的最大回撤率交易日
        min_point_df = stock_dat.sort_values(by = ['per_close'])[0:1]
        min_point_close = min_point_df.per_close
        #找出收盘价最高值交易日，并打印显示
        max_point_df = stock_dat[stock_dat.index <= min_point_close.index[0]].sort_values(by = ['Close'],ascending=False)[0:1]
        max_point_close = max_point_df.Close
        
        #打印收盘价的最大回撤率与所对应的最高值交易日和最大回撤交易日
        print(u"股价最大回撤%5.2f%% 从%s开始至%s结束" % ((1 - min_point_close.values[0]) * 100, max_point_close.index[0], min_point_close.index[0]))

        type_dict = {'最大收盘价': stock_dat.max_close, "收盘价": stock_dat.Close}
        
        view_function = MplTypesDraw.mpl.route_output(u"line")
        view_function(stock_dat.index, type_dict, sub_graph)

        type_dict = {
            u"股价最大回撤\n{}".format(1 - min_point_close.values):
                {
                    'andata': min_point_df,
                    'va':'top',
                    'xy_y': 'Close',
                    'xytext':(0, stock_dat['High'].mean()),
                    'fontsize': 8,
                    'arrow':dict(facecolor='green', shrink =  0.1)
                },
        }
        view_function = MplTypesDraw.mpl.route_output(u"annotate")
        view_function(stock_dat.index, type_dict, sub_graph)

    @app.route_types(u"trade")
    def trade_graph(stock_dat, sub_graph, para_dat):
        # 交易获利/亏损区间可视化
        type_dict = {
            'signal': stock_dat.Signal,
            'jdval': stock_dat.Close,
            'va':'top',
            'xy_y':'Close',
            'xytext':(0, stock_dat['High'].mean()),
            'fontsize': 0,
            'arrow': dict(facecolor='yellow',shrink=0.1)
        }
        view_function = MplTypesDraw.mpl.route_output(u"filltrade")
        view_function(stock_dat.index, type_dict, sub_graph)

    def __init__(self, **kwargs):
        MplTypesDraw.__init__(self)
        # 创建fig对象
        self.fig = plt.figure(figsize=(kwargs['figsize']), dpi=100, facecolor="white")
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False 
        self.graph_dict = {}
        self.graph_curr = {}
        try:
            gs = gridspec.GridSpec(kwargs['nrows'], kwargs['ncols'],left = kwargs['left'],bottom=kwargs['bottom'],right=kwargs['right'],top=kwargs['top'],wspace = kwargs['wspace'],
                                   hspace=kwargs['hspace'],height_ratios=kwargs['height_ratios'])
        except:
            raise Exception("para error")
        else:
            for i in range(0, kwargs['nrows'], 1):
                self.graph_dict[kwargs['subplots'][i]] = self.fig.add_subplot(gs[i,:])
    
    def log_trade_info(self, stock_dat):
        signal_shift = stock_dat.Signal.shift(1)
        signal_shift.fillna(value=-1, inplace=True) #序列最前面的NaN值用-1填充
        list_signal = np.sign(stock_dat.Signal - signal_shift)
        
        buy_signal = stock_dat[list_signal.isin([1])]
        sell_signal = stock_dat[list_signal.isin([-1])]
        
        if len(buy_signal) > len(sell_signal):
            buy_signal = buy_signal.drop(buy_signal.index[-1])

        trade_info = pd.DataFrame({'BuyTime':buy_signal.index.strftime("%y.%m.%d"),
                                  'SellTime':sell_signal.index.strftime("%y.%m.%d"),
                                  'BuyPrice':buy_signal.Close.values,
                                  'SellPrice':sell_signal.Close.values})
        
        trade_info['DiffPrice'] = trade_info.SellPrice - trade_info.BuyPrice
        trade_info['PctProfit'] = np.round(trade_info.DiffPrice / trade_info.BuyPrice * 100,2)

        win_count = (trade_info.DiffPrice >= 0).sum()
        loss_count = (trade_info.DiffPrice <0).sum()
        win_profit = trade_info[trade_info.PctProfit >= 0].PctProfit.sum()
        loss_profit = trade_info[trade_info.PctProfit < 0].PctProfit.sum()

        print(trade_info)
        print(f'亏损次数:{loss_count}, 盈利次数:{win_count}, 胜率:{round(win_count / (win_count +loss_count) * 100,2)}%')
        if loss_count == 0:
            loss_count = 1
        print(f'平均亏损:{round((loss_profit / loss_count),2)} % 平均盈利:{round((win_profit / win_count),2)}%')
    
    def graph_run(self, stock_data, **kwargs):
        #绘制子图
        self.df_ohlc = stock_data
        # 临时把标准输出出重定向到一个文件，然后再反复正常
        with open('logtrade.txt', 'w',encoding='utf-8') as f:
            oldstdout = sys.stdout
            sys.stdout = f
            try:
                self.log_trade_info(self.df_ohlc)
                for key in kwargs:
                    self.graph_curr = self.graph_dict[kwargs[key]['graph_name']]
                    for path, val in kwargs[key]['graph_type'].items():
                        view_function = MultiTraceIf.app.route_output(path)
                        view_function(self.df_ohlc, self.graph_curr, val)
                    self.graph_attr(**kwargs[key])
                plt.show()
            finally:
                sys.stdout = oldstdout
    
    def graph_attr(self, **kwargs):
        if 'title' in kwargs.keys():
            self.graph_curr.set_title(kwargs['title'])
        
        if 'legend' in kwargs.keys():
            self.graph_curr.legend(loc=kwargs['legend'], shadow=True)
        
        if 'xlabel' in kwargs.keys():
            self.graph_curr.set_xlabel(kwargs['xlabel'])

        self.graph_curr.set_ylabel(kwargs['ylabel'])
        self.graph_curr.set_xlim(0, len(self.df_ohlc.index)) # 设置一下x轴的范围
        self.graph_curr.set_xticks(range(0, len(self.df_ohlc.index), kwargs['xticks'])) # x轴刻度设定 每15天标一个日期

        if 'xticklabels' in kwargs.keys():
            self.graph_curr.set_xticklabels(
                [self.df_ohlc.index.strftime(kwargs['xticklabels'])[index] for index in
                 self.graph_curr.get_xticks()]) # 标签设置为日期
            
            # x轴每个ticker标签都向右45度
            for label in self.graph_curr.xaxis.get_ticklabels():
                label.set_rotation(45)
                label.set_fontsize(10) #设置标签字体
        else:
            for label in self.graph_curr.xaxis.get_ticklabels():
                label.set_visible(False)
            


# 登陆系统
def bs_k_data_stock(code_val='sz.000651',start_val = '2025-01-01', end_val='2025-05-01',freq_val='d',adjust_val='3'):
    lg = bs.login()
    # 获取历史行情数据
    fields = "date,open,high,low,close,volume"
    df_bs = bs.query_history_k_data_plus(code_val, fields, start_date=start_val, end_date=end_val,frequency=freq_val ,adjustflag=adjust_val)
    #日K线，3，默认不赋权，1 后复权，2 前复权
    data_list= []
    while (df_bs.error_code == '0') & df_bs.next():
        #获取一条记录，将记录合并在一起
        data_list.append(df_bs.get_row_data())
    result = pd.DataFrame(data_list, columns=df_bs.fields)

    result.close = result.close.astype('float64')
    result.open = result.open.astype('float64')
    result.low = result.low.astype('float64')
    result.high = result.high.astype('float64')
    result.volume = result.volume.astype('float64')
    result.volume = result.volume/100 #单位转换:股-手
    result.date = pd.DatetimeIndex(result.date)
    result.set_index("date", drop=True, inplace= True)
    result.index = result.index.set_names('Date')

    recon_data = {'High' : result.high, 'Low': result.low, 'Open': result.open, 'Close':result.close, 'Volume': result.volume}
    df_recon = pd.DataFrame(recon_data)

    #退出系统
    bs.logout()
    return df_recon


def get_trade_signal(stock_dat):
    df_csvload_trade = pd.read_csv('601595.csv',index_col=None,parse_dates=[2,3])
    stock_dat = stock_dat.assign(Signal = np.nan)
    stock_dat.loc[df_csvload_trade["Buy-Time"],'Signal'] = 1
    stock_dat.loc[df_csvload_trade["Sell-Time"],'Signal'] = -1
    # 与前面元素保持一致
    stock_dat.ffill(inplace=True)
    # 序列最前面几个NaN用-1填充
    stock_dat.fillna(value = -1, inplace = True)
    return stock_dat

def get_ndays_signal(stock_dat, N1= 15, N2 = 5):
    # 海龟策略，唐奇安通道突破（N日突破）买入/卖出信号
    stock_dat['N1_High'] = stock_dat.High.rolling(window=N1).max() #计算最近N1个交易日最高价
    expan_max = stock_dat.High.expanding().max() #滚动计算当前交易日为止的最大值
    stock_dat.fillna({'N1_High':expan_max},inplace=True) #填充前N1个nan
    stock_dat['N2_Low'] = stock_dat.Low.rolling(window=N2).min() #计算最近N2个交易日最低价
    expan_min = stock_dat.Low.expanding().min()
    stock_dat.fillna({'N2_Low':expan_min},inplace=True) #目前出现过的最小值填充前N2个nan
    # 收盘价超过N1最高价，买入股票
    buy_index = stock_dat[stock_dat.Close > stock_dat.N1_High.shift(1)].index
    # 收盘价超过N2最高价，卖出股票
    sell_index = stock_dat[stock_dat.Close < stock_dat.N2_Low.shift(1)].index
    stock_dat.loc[buy_index, 'Signal'] = 1
    stock_dat.loc[sell_index,'Signal'] = -1
    stock_dat['Signal'] = stock_dat.Signal.shift(1)
    stock_dat.ffill( inplace= True) #与前面元素保持一致
    stock_dat.fillna({'Signal':-1},inplace = True) #序列最前面几个NaN值用-1填充
    return stock_dat


def get_ndays_atr_signal(stock_dat, N1 = 15, N2 = 5, n_win = 3.5, n_loss = 1.8):
    # 海龟策略-唐奇安通道突破（N日突破）买入/卖出信号
    stock_dat['N1_High'] = stock_dat.High.rolling(window=N1).max() #计算最近N1个交易日最高价
    expan_max = stock_dat.High.expanding().max() #滚动计算当前交易日为止的最大值
    stock_dat.fillna({'N1_High':expan_max}, inplace=True) #填充前N1个nan
    stock_dat['N2_Low'] = stock_dat.Low.rolling(window=N2).min() #计算租金N2个交易日最低价
    expan_min = stock_dat.Low.expanding().min()
    stock_dat.fillna({'N2_Low':expan_min}, inplace=True) #目前出现过的最小值填充前N2个nan

    stock_dat['ATR21'] = talib.ATR(stock_dat.High.values, stock_dat.Low.values, stock_dat.Close.values, timeperiod=21)
    # 收盘价超过N1最高价 买入股票
    buy_index = stock_dat[stock_dat.Close > stock_dat.N1_High.shift(1)].index
    # 收盘价超过N2最低价 卖出股票
    sell_index = stock_dat[stock_dat.Close < stock_dat.N2_Low.shift(1)].index
    stock_dat.loc[buy_index, 'Signal'] = 1
    stock_dat.loc[sell_index, 'Signal'] = -1
    stock_dat['Signal'] = stock_dat.Signal.shift(1)
    buy_price = 0

    for k1_index, today in stock_dat.iterrows():
        if (buy_price == 0) and (today.Signal == 1):
            buy_price = today.Close
        # 到达收盘价少于买入价后触发卖出
        elif(buy_price != 0) and (buy_price > today.Close) and ((buy_price-today.Close) > n_loss * today.ATR21):
            print(f'止损时间:{k1_index.strftime("%y.%m.%d")} and 止损价格:{round(today.Close,2)}')
            stock_dat.loc[k1_index, 'Signal'] = -1
        # 到达收盘价多余买入价后触发卖出
        elif(buy_price  != 0) and (buy_price < today.Close) and ((today.Close - buy_price) > n_win * today.ATR21):
            print(f'止盈时间:{k1_index.strftime("%y.%m.%d")} and 止盈价格:{round(today.Close,2)}')
            stock_dat.loc[k1_index, 'Signal'] = -1
            buy_price = 0 
        else:
            pass
    stock_dat['Signal'] = stock_dat['Signal'].ffill() #与前面元素保持一致
    stock_dat.fillna({'Signal':-1}, inplace=True) #序列最前面几个NaN值用0填充
    return stock_dat
        


df_stockload = bs_k_data_stock('sh.600111','2025-01-01', '2025-08-08')
layout_dict = {
    'figsize':(14, 8),
    'nrows': 3,
    'ncols': 1,
    'left': 0.08,
    'bottom':0.15,
    'right':0.95,
    'top':0.95,
    'wspace':None,
    'hspace':0,
    'height_ratios': [1.5, 1, 1],
    'subplots':['kgraph','cashgraph','cmppfgraph']
}
subplots_dict = {'graph_fst':{'graph_name':'kgraph',
                              'graph_type':{'trade':None,
                                            'close_retrace':None
                                            },
                              'title':u"回测分析",
                              'ylabel':u'价格',
                              'xticks':15,
                              'legend':'best'
                              },
                'graph_sec':{'graph_name':'cashgraph',
                              'graph_type':{'cash_profit':100000 #初始资金
                                            },
                              'ylabel':u'资金收益和回撤',
                              'xticks':15,
                              'legend':'best'
                              },
                'graph_fth':{'graph_name':'cmppfgraph',
                              'graph_type':{'cmp_profit':None
                                            },
                              'ylabel':u'策略收益vs基准收益',
                              'xlabel':u'日期',
                              'xticks':15,
                              'legend':'best',
                              'xticklabel':'%Y-%m-%d'
                              },
                              }
draw_stock = MultiTraceIf(**layout_dict)
df_stockload2 = get_ndays_atr_signal(df_stockload)
#df_stockload2 = get_trade_signal(df_stockload)
#df_stockload2 = get_ndays_signal(df_stockload)
draw_stock.graph_run(df_stockload2, **subplots_dict)









